#include "CaeConnectionPool.h"

#include <utility>

CaeConnectionPool::CaeConnectionPool() = default;

CaeConnectionPool::~CaeConnectionPool() {
    Clear();
}

void CaeConnectionPool::Add(std::string connectionId, std::shared_ptr<CaeSocketAgent> socketAgent)
{
    std::unique_lock<std::shared_mutex> lock(m_mapMutex);
    m_socketAgentMap[connectionId] = std::move(socketAgent);
}

void CaeConnectionPool::Remove(const std::shared_ptr<CaeSocketAgent>& socketAgent)
{
    std::unique_lock<std::shared_mutex> lock(m_mapMutex);
    for (auto it = m_socketAgentMap.begin(); it != m_socketAgentMap.end(); it++) {
        if (it->second == socketAgent) {
            it->second = nullptr;
            m_socketAgentMap.erase(it);
            break;
        }
    }
}

void CaeConnectionPool::Remove(std::string &connectionId)
{
    std::unique_lock<std::shared_mutex> lock(m_mapMutex);
    auto iter = m_socketAgentMap.find(connectionId);
    if(iter != m_socketAgentMap.end()) {
        iter->second = nullptr;
        m_socketAgentMap.erase(iter);
    }
}

void CaeConnectionPool::Clear()
{
    std::unique_lock<std::shared_mutex> lock(m_mapMutex);
    m_socketAgentMap.clear();
}

size_t CaeConnectionPool::Size()
{
    std::shared_lock<std::shared_mutex> lock(m_mapMutex);
    return m_socketAgentMap.size();
}

std::vector<std::shared_ptr<CaeSocketAgent>> CaeConnectionPool::GetAll()
{
    std::vector<std::shared_ptr<CaeSocketAgent>> all;
    std::shared_lock<std::shared_mutex> lock(m_mapMutex);
    for (const auto& pair : m_socketAgentMap) {
        all.push_back(pair.second);
    }
    return all;
}

std::vector<std::shared_ptr<CaeSocketAgent>> CaeConnectionPool::GetAllValid()
{
    std::vector<std::shared_ptr<CaeSocketAgent>> all;
    std::shared_lock<std::shared_mutex> lock(m_mapMutex);
    for (const auto& pair : m_socketAgentMap) {
        if (pair.second->IsVerifyComplete() && pair.second->GetSocketAgentState() == CAE::RUNNING) {
            all.push_back(pair.second);
        }
    }

    return all;
}

std::shared_ptr<CaeSocketAgent> CaeConnectionPool::GetIfExist(const std::string& connectionId)
{
    std::shared_lock<std::shared_mutex> lock(m_mapMutex);
    if (m_socketAgentMap.count(connectionId) != 0) {
        return m_socketAgentMap[connectionId];
    } else {
        return nullptr;
    }
}

std::shared_ptr<CaeSocketAgent> CaeConnectionPool::Pop(const std::string& connectionId)
{
    std::shared_ptr<CaeSocketAgent> pvalue = nullptr;
    std::unique_lock<std::shared_mutex> lock(m_mapMutex);
    auto iter = m_socketAgentMap.find(connectionId);
    if(iter != m_socketAgentMap.end()) {
        pvalue = m_socketAgentMap[connectionId];
        m_socketAgentMap.erase(iter);
    }

    return pvalue;
}
